#include <atomic>
#include <thread>
#include <vector>
#include <mutex>
#include <memory>
#include <iostream>

#include <boost/pool/object_pool.hpp>
#include <boost/asio.hpp>
#include <boost/lockfree/queue.hpp>


namespace asio = boost::asio;
using asio::ip::tcp;

// 跨平台编译宏
#if defined(_WIN32)
    #define IOCP_MODE // Windows 使用 IOCP
#else
    #define EPOLL_MODE // Linux 使用 epoll
    #include <pthread.h>
#endif

// --------------- 核心配置 ---------------
constexpr int MAX_CONNECTIONS = 100000; // 最大连接数
constexpr int BUFFER_SIZE = 1024;       // 数据包大小 (1KB)
constexpr int HEARTBEAT_INTERVAL = 30;   // 心跳间隔(秒)

// --------------- 协议处理接口 ---------------
class IProtocolHandler {
public:
    virtual ~IProtocolHandler() = default;
    virtual void parse_request(uint8_t* data, size_t len) = 0;
    virtual void build_response(uint8_t* data, size_t len) = 0;
};

// --------------- 数据结构 ---------------
class DeviceSession {
public:
    DeviceSession(asio::io_context& io_ctx) 
        : socket(io_ctx),
        heartbeat_timer(io_ctx),
        buffer(), 
        is_active(true) {}

    ~DeviceSession() {
    if (socket.is_open()) {
        boost::system::error_code ec;
        socket.cancel();
        socket.shutdown(tcp::socket::shutdown_both, ec); // 安全关闭
        socket.close(ec);
    }
}

    tcp::socket socket;
    asio::steady_timer heartbeat_timer;
    std::array<uint8_t, BUFFER_SIZE> buffer;
    std::atomic<bool> is_active;
};

// --------------- 连接池管理 ---------------
class ConnectionPool {
    boost::lockfree::queue<DeviceSession*> pool_queue_{MAX_CONNECTIONS};  // 无锁队列
    std::atomic<size_t> active_count_{0};  // 原子计数器统计活跃连接数

public:
    DeviceSession* acquire(asio::io_context& io_ctx) {
        DeviceSession* session = nullptr;
        // 尝试从队列弹出可用连接
        if (pool_queue_.pop(session)) {
            session->is_active = true;
            return session;
        }
        // 无可用连接时新建（不超过上限）
        if (active_count_.load() < MAX_CONNECTIONS) {
            session = new DeviceSession(io_ctx);
            active_count_.fetch_add(1);  // 原子递增
            session->is_active = true;
            return session;
        }
        return nullptr;  // 连接数已达上限
    }

    void release(DeviceSession* session) {
        if (session) {
            session->is_active = false;
            pool_queue_.push(session);  // 归还到无锁队列（非销毁）
        }
    }

    ~ConnectionPool() {
        // 1. 清空队列中的空闲连接
        DeviceSession* session = nullptr;
        while (pool_queue_.pop(session)) {
            delete session;
            active_count_.fetch_sub(1); // 更新计数器
        }
        
        // 2. 检查是否有未归还的活跃连接（可选）
        if (active_count_.load() > 0) {
            // 记录警告日志或采取恢复措施
        }
    }
};

// --------------- 心跳管理器 ---------------
class HeartbeatManager {
public:
    static void send_heartbeat(DeviceSession* session, 
                              std::shared_ptr<IProtocolHandler> handler) {
        if (!session || !session->is_active) return;

        // 构造心跳包 (Modbus功能码0x0B)
        uint8_t heartbeat[] = {0x00, 0x01, 0x00, 0x00, 0x00, 0x06, 0x01, 0x0B};
        
        asio::async_write(
            session->socket, 
            asio::buffer(heartbeat),
            [session, handler](const boost::system::error_code& ec, size_t /*bytes*/) {
                if (ec) {
                    if (!session || !session->is_active) return;
                    session->is_active = false;
                    std::cerr << "Heartbeat failed: " << ec.message() << std::endl;
                }
            }
        );
    }
};

// --------------- I/O 上下文线程 ---------------
class IOContextThread {
    asio::io_context io_ctx_;
    std::unique_ptr<asio::io_context::work> work_;
    std::thread thread_;
    ConnectionPool& pool_;
    std::shared_ptr<IProtocolHandler> protocol_handler_;

public:
    IOContextThread(ConnectionPool& pool, std::shared_ptr<IProtocolHandler> handler)
        : work_(std::make_unique<asio::io_context::work>(io_ctx_)),
          pool_(pool),
          protocol_handler_(handler) {
        thread_ = std::thread([this] { 
            #if !defined(_WIN32)
                // Linux线程绑定CPU核心
                cpu_set_t cpuset;
                CPU_ZERO(&cpuset);
                size_t core_id = std::hash<std::thread::id>{}(std::this_thread::get_id()) % 
                                 std::thread::hardware_concurrency();
                CPU_SET(core_id, &cpuset);
                pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
            #endif
            io_ctx_.run(); 
        });
    }

    void start_accept(tcp::acceptor& acceptor) {
        auto* session = pool_.acquire(io_ctx_);
        if (!session) return;

        acceptor.async_accept(
            session->socket,
            [this, session, &acceptor](const boost::system::error_code& ec) {
                if (!ec && session->is_active) {
                    start_async_read(session);
                    start_heartbeat(session);
                } else if (ec) {
                    std::cerr << "Accept error: " << ec.message() << std::endl;
                    pool_.release(session);
                }
                start_accept(acceptor);
            }
        );
    }

private:
    void start_async_read(DeviceSession* session) {
        session->socket.async_read_some(
            asio::buffer(session->buffer),
            [this, session](const boost::system::error_code& ec, size_t bytes) {
                if (ec || !session->is_active) return handle_io_error(session, ec);
                if (!protocol_handler_) return;
                protocol_handler_->parse_request(session->buffer.data(), bytes);
                send_response(session);
                start_async_read(session);
                
            }
        );
    }

    void send_response(DeviceSession* session) {
        protocol_handler_->build_response(session->buffer.data(), session->buffer.size());
        asio::async_write(
            session->socket,
            asio::buffer(session->buffer),
            [session](const boost::system::error_code& ec, size_t /*bytes*/) {
                if (ec) {
                    if (!session || !session->is_active) return;

                    session->is_active = false;
                    std::cerr << "Write error: " << ec.message() << std::endl;
                }
            }
        );
    }

    void start_heartbeat(DeviceSession* session) {
        session->heartbeat_timer.expires_after(std::chrono::milliseconds(HEARTBEAT_INTERVAL*1000));
        
        session->heartbeat_timer.async_wait(
            [this, session](const boost::system::error_code& ec) {
                try{
                    if (!ec && session->is_active) {
                    HeartbeatManager::send_heartbeat(session, protocol_handler_);
                    start_heartbeat(session);
                    } else {
                        handle_io_error(session, ec);
                    }
                }
                catch (const std::exception& e) {
                    std::cerr << "Unhandled exception: " << e.what();
                    session->is_active = false;
                }
            }
        );
    }

    void handle_io_error(DeviceSession* session, const boost::system::error_code& ec) {
        if (ec) {
            std::cerr << "I/O error: " << ec.message() << std::endl;
        }
        session->is_active = false;
        pool_.release(session);
    }
};

// --------------- 反应器主线程 ---------------
class Reactor {
    asio::io_context main_ctx_;
    tcp::acceptor acceptor_;
    std::vector<std::unique_ptr<IOContextThread>> io_threads_;
    ConnectionPool connection_pool_;
    std::shared_ptr<IProtocolHandler> protocol_handler_;
    std::thread main_thread_;

public:
    Reactor(uint16_t port, std::shared_ptr<IProtocolHandler> handler)
        : acceptor_(main_ctx_, tcp::endpoint(tcp::v4(), port)),
          protocol_handler_(handler) {
        
        // 初始化I/O线程池
        const size_t thread_count = std::thread::hardware_concurrency();
        for (size_t i = 0; i < thread_count; ++i) {
            io_threads_.emplace_back(
                std::make_unique<IOContextThread>(connection_pool_, protocol_handler_));
        }
        
        // 启动连接分发
        start_accept_loop();
        
        // 运行主线程
        main_thread_ = std::thread([this] { main_ctx_.run(); });
    }

private:
    void start_accept_loop() {
        static std::atomic<size_t> next_thread = 0;
        
        acceptor_.async_accept(
            [this](const boost::system::error_code& ec, tcp::socket socket) {
                if (!ec) {
                    size_t idx = next_thread++ % io_threads_.size(); // 原子递增
                    io_threads_[idx]->start_accept(acceptor_);
                } else {
                    std::cerr << "Accept error: " << ec.message() << std::endl;
                }
                start_accept_loop();
            }
        );
    }
};

// --------------- 平台初始化 ---------------
void init_network() {
    #if defined(_WIN32)
        WSADATA wsaData;
        if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
            throw std::runtime_error("WSAStartup failed");
        }
    #endif
}

// --------------- Modbus协议处理器 ---------------
class ModbusHandler : public IProtocolHandler {
public:
    void parse_request(uint8_t* data, size_t len) override {
        // 实际Modbus协议解析实现
        if (len > 7) {
            uint8_t func_code = data[7];
            std::cout << "Received Modbus function: 0x" << std::hex 
                      << static_cast<int>(func_code) << std::endl;
        }
    }

    void build_response(uint8_t* data, size_t len) override {
        // 实际Modbus响应构造实现
        if (len > 0) data[0] = 0x01; // 示例：设置事务ID
    }
};

// --------------- 主函数 ---------------
int main() {
    init_network();
    
    // 创建协议处理器
    auto handler = std::make_shared<ModbusHandler>();
    
    // 启动服务器
    Reactor server(5020, handler);
    
    // 主线程等待
    std::this_thread::sleep_until(std::chrono::steady_clock::time_point::max());
    return 0;
}